import numpy as np
import os
import torchvision
from PIL import Image
import torch.utils.data as data_utils
from typing import List, Tuple, Callable, Union
from .utils import xml2dict

__all__ = ["XMLReader"]


class DetectionDataset(data_utils.Dataset):
    """
    samples:List[str], 图片的路径数组
    boxs: List[List[List[int]]], 边界框数组
    labels: List[List[int]], 类别数组
    """

    def __init__(self, samples, boxs, labels, transforms=None):
        super().__init__()
        assert len(samples) == len(boxs) == len(labels)
        assert isinstance(samples, (list, tuple))
        assert isinstance(boxs, (list, tuple))
        assert isinstance(labels, (list, tuple))
        self.samples = samples
        self.boxs = boxs
        self.labels = labels
        self.transforms = transforms

    def __getitem__(self, idx):
        image = Image.open(self.samples[idx]).convert("RGB")
        # noinspection PyTypeChecker
        image = np.array(image, dtype=np.float32)
        bboxes = self.boxs[idx]
        labels = self.labels[idx]
        if self.transforms:
            o = self.transforms(image=image, bboxes=bboxes, labels=labels)
            image, bboxes, labels = o["image"], o["bboxes"], o["labels"]
        target = {"boxes": bboxes, "labels": labels}
        return image, target

    def __len__(self):
        return len(self.samples)


class XMLReader(DetectionDataset):
    def __init__(self, img_dir, xml_dir, txt_name,
                 img_ext=".jpg", transfrom=None, start_label_id=0):
        self.class_to_idx = {}
        self.start_label_id = start_label_id
        samples, boxs, labels = self.process(img_dir, xml_dir, txt_name)
        self.idx_to_class = {v: k for k, v in self.class_to_idx.items()}
        super().__init__(samples, boxs, labels, transfrom)

    @property
    def num_classes(self):
        return len(self.class_to_idx)

    def process(self, img_dir, xml_dir, csv_name, img_ext=".jpg"):
        assert os.path.exists(csv_name), f"{csv_name} 不存在"
        if isinstance(csv_name, str):
            with open(csv_name, "r", encoding="utf8") as f:
                name_list = [line[:-1] for line in f]

        samples, boxes, labels = [], [], []
        cat = lambda dir, name, ext: os.path.join(dir, name + ext)
        for name in name_list:
            objects = xml2dict(cat(xml_dir, name, ".xml"))["annotation"]
            if "object" not in objects:
                continue  # 如果图片内没有物体，就不读取
            samples.append(cat(img_dir, name, img_ext))
            objects = objects["object"]
            objects = objects if isinstance(objects, list) else [objects]
            box = []
            label = []
            for obj in objects:
                label_name = obj["name"]
                if label_name in self.class_to_idx:
                    label_id = self.class_to_idx[label_name]
                else:
                    label_id = len(self.class_to_idx) + self.start_label_id
                    self.class_to_idx[label_name] = label_id
                x1 = int(obj["bndbox"]["xmin"])
                y1 = int(obj["bndbox"]["ymin"])
                x2 = int(obj["bndbox"]["xmax"])
                y2 = int(obj["bndbox"]["ymax"])
                label.append(label_id)
                box.append([x1, y1, x2, y2])
            boxes.append(box)
            labels.append(label)
        return samples, boxes, labels

    def get_dataloader(self, batch_size=32, shuffle=False, collate_fn=lambda batch: zip(*batch)):
        return data_utils.DataLoader(self, batch_size=batch_size, shuffle=shuffle, collate_fn=collate_fn)

    def export_yolov5_dataset(self, root):
        import yaml
        create_dir_if_not_exists = lambda path: os.makedirs(path) if not os.path.exists(path) else None
        create_dir_if_not_exists(root)
        create_dir_if_not_exists(os.path.join(root, "images"))
        create_dir_if_not_exists(os.path.join(root, "labels"))

        for i, (sample, box, label) in enumerate(zip(self.samples, self.boxs, self.labels)):
            if os.path.exists(os.path.join(root, "images", f"{i}.jpg")):
                continue
            img = Image.open(sample).convert("RGB")
            img.save(os.path.join(root, "images", f"{i}.jpg"))
            with open(os.path.join(root, "labels", f"{i}.txt"), "w", encoding="utf8") as f:
                for b, l in zip(box, label):
                    x1, y1, x2, y2 = b
                    cx, cy, w, h = (x1 + x2) / 2, (y1 + y2) / 2, x2 - x1, y2 - y1
                    cx, cy, w, h = cx / img.width, cy / img.height, w / img.width, h / img.height
                    f.write(f"{l} {cx} {cy} {w} {h}\n")
        # 创建data.yaml
        yaml_data = {
            "path": root,
            "train": f"{root}/images",
            "val": f"{root}/images",
            "nc": self.num_classes,
            "names": self.idx_to_class}
        yaml.dump(yaml_data, open(os.path.join(root, "data.yaml"), "w", encoding="utf8"), allow_unicode=True)
